from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
import sys
import os

# 添加项目路径到 Python 路径
sys.path.append('/app')

from app.services.etl_service import get_etl_service
from app.utils.doris_client import get_doris_client
from app.database import get_db

# DAG 默认参数
default_args = {
    'owner': 'dataworks',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}

# 创建 DAG
dag = DAG(
    'doris_etl_dag',
    default_args=default_args,
    description='Doris 数据仓库 ETL 数据同步 DAG',
    schedule_interval=timedelta(hours=1),  # 每小时执行一次
    catchup=False,
    tags=['etl', 'doris', 'datawarehouse']
)


def check_doris_connection(**context):
    """检查 Doris 连接"""
    print("正在检查 Doris 数据仓库连接...")
    
    doris_client = get_doris_client()
    is_connected = doris_client.test_connection()
    
    if is_connected:
        print("✅ Doris 连接测试成功")
        
        # 获取数据库信息
        tables = doris_client.get_tables()
        print(f"📊 Doris 数据库 '{doris_client.database}' 中有 {len(tables)} 个表")
        
        for table in tables[:5]:  # 显示前5个表
            print(f"  - {table.get('table_name', 'N/A')}: {table.get('table_comment', '无描述')}")
            
        return True
    else:
        print("❌ Doris 连接测试失败")
        raise Exception("无法连接到 Doris 数据仓库")


def sync_sample_data(**context):
    """同步样例数据到 Doris"""
    print("开始同步样例数据到 Doris...")
    
    doris_client = get_doris_client()
    
    # 创建样例表（如果不存在）
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS sample_sales (
        id BIGINT,
        order_date DATE,
        customer_id INT,
        product_name STRING,
        quantity INT,
        price DECIMAL(10,2),
        total_amount DECIMAL(10,2),
        region STRING,
        sales_rep STRING,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    ) ENGINE=OLAP
    DUPLICATE KEY(id, order_date)
    DISTRIBUTED BY HASH(id) BUCKETS 8
    PROPERTIES (
        "replication_num" = "1"
    )
    """
    
    try:
        doris_client.execute_sql(create_table_sql)
        print("✅ 样例表 'sample_sales' 创建成功")
    except Exception as e:
        print(f"⚠️ 表可能已存在: {str(e)}")
    
    # 插入样例数据
    sample_data = [
        {
            'id': 1001,
            'order_date': '2024-01-15',
            'customer_id': 101,
            'product_name': '智能手机',
            'quantity': 2,
            'price': 2999.00,
            'total_amount': 5998.00,
            'region': '华东',
            'sales_rep': '张三'
        },
        {
            'id': 1002,
            'order_date': '2024-01-15',
            'customer_id': 102,
            'product_name': '笔记本电脑',
            'quantity': 1,
            'price': 6999.00,
            'total_amount': 6999.00,
            'region': '华南',
            'sales_rep': '李四'
        },
        {
            'id': 1003,
            'order_date': '2024-01-16',
            'customer_id': 103,
            'product_name': '平板电脑',
            'quantity': 3,
            'price': 1999.00,
            'total_amount': 5997.00,
            'region': '华北',
            'sales_rep': '王五'
        }
    ]
    
    success = doris_client.bulk_insert('sample_sales', sample_data)
    
    if success:
        print(f"✅ 成功插入 {len(sample_data)} 条样例数据")
        
        # 验证数据
        result = doris_client.execute_query("SELECT COUNT(*) as total FROM sample_sales")
        total_count = result[0]['total'] if result else 0
        print(f"📊 表 'sample_sales' 当前总记录数: {total_count}")
        
    else:
        raise Exception("样例数据插入失败")


def generate_daily_report(**context):
    """生成日报数据"""
    print("开始生成日报数据...")
    
    doris_client = get_doris_client()
    
    # 创建日报汇总表
    create_report_table_sql = """
    CREATE TABLE IF NOT EXISTS daily_sales_report (
        report_date DATE,
        region STRING,
        total_orders INT,
        total_quantity INT,
        total_amount DECIMAL(15,2),
        avg_order_amount DECIMAL(10,2),
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    ) ENGINE=OLAP
    DUPLICATE KEY(report_date, region)
    DISTRIBUTED BY HASH(report_date) BUCKETS 4
    PROPERTIES (
        "replication_num" = "1"
    )
    """
    
    try:
        doris_client.execute_sql(create_report_table_sql)
        print("✅ 日报表 'daily_sales_report' 创建成功")
    except Exception as e:
        print(f"⚠️ 表可能已存在: {str(e)}")
    
    # 生成今日报告数据
    today = context['ds']  # Airflow 提供的日期参数
    
    report_sql = f"""
    INSERT INTO daily_sales_report (report_date, region, total_orders, total_quantity, total_amount, avg_order_amount)
    SELECT 
        '{today}' as report_date,
        region,
        COUNT(*) as total_orders,
        SUM(quantity) as total_quantity,
        SUM(total_amount) as total_amount,
        AVG(total_amount) as avg_order_amount
    FROM sample_sales 
    WHERE order_date = '{today}'
    GROUP BY region
    """
    
    try:
        doris_client.execute_sql(report_sql)
        print(f"✅ {today} 日报数据生成成功")
        
        # 查看生成的报告
        result = doris_client.execute_query(f"""
            SELECT * FROM daily_sales_report 
            WHERE report_date = '{today}'
            ORDER BY total_amount DESC
        """)
        
        print(f"📊 {today} 日报数据:")
        for row in result:
            print(f"  - {row['region']}: {row['total_orders']}单, "
                  f"数量{row['total_quantity']}, 金额{row['total_amount']}")
            
    except Exception as e:
        print(f"❌ 日报数据生成失败: {str(e)}")
        raise


def data_quality_check(**context):
    """数据质量检查"""
    print("开始数据质量检查...")
    
    doris_client = get_doris_client()
    
    # 检查空值
    null_check_sql = """
    SELECT 
        'sample_sales' as table_name,
        SUM(CASE WHEN customer_id IS NULL THEN 1 ELSE 0 END) as null_customer_id,
        SUM(CASE WHEN product_name IS NULL OR product_name = '' THEN 1 ELSE 0 END) as null_product_name,
        SUM(CASE WHEN quantity <= 0 THEN 1 ELSE 0 END) as invalid_quantity,
        SUM(CASE WHEN price <= 0 THEN 1 ELSE 0 END) as invalid_price,
        COUNT(*) as total_records
    FROM sample_sales
    """
    
    result = doris_client.execute_query(null_check_sql)
    if result:
        quality_report = result[0]
        print("📋 数据质量检查报告:")
        print(f"  总记录数: {quality_report['total_records']}")
        print(f"  空客户ID: {quality_report['null_customer_id']}")
        print(f"  空产品名: {quality_report['null_product_name']}")
        print(f"  无效数量: {quality_report['invalid_quantity']}")
        print(f"  无效价格: {quality_report['invalid_price']}")
        
        # 计算数据质量得分
        total_issues = (quality_report['null_customer_id'] + 
                       quality_report['null_product_name'] + 
                       quality_report['invalid_quantity'] + 
                       quality_report['invalid_price'])
        
        quality_score = 100 - (total_issues / quality_report['total_records'] * 100) if quality_report['total_records'] > 0 else 0
        print(f"  数据质量得分: {quality_score:.2f}%")
        
        if quality_score < 95:
            print("⚠️ 数据质量警告: 得分低于95%")
        else:
            print("✅ 数据质量良好")


def optimize_tables(**context):
    """优化表性能"""
    print("开始优化表性能...")
    
    doris_client = get_doris_client()
    
    tables_to_optimize = ['sample_sales', 'daily_sales_report']
    
    for table in tables_to_optimize:
        try:
            # 获取表统计信息
            stats = doris_client.get_table_stats(table)
            print(f"📊 表 {table} 统计信息: {stats}")
            
            # Doris 表优化
            success = doris_client.optimize_table(table)
            if success:
                print(f"✅ 表 {table} 优化成功")
            else:
                print(f"⚠️ 表 {table} 优化失败")
                
        except Exception as e:
            print(f"❌ 优化表 {table} 时出错: {str(e)}")


# 定义任务
check_connection_task = PythonOperator(
    task_id='check_doris_connection',
    python_callable=check_doris_connection,
    dag=dag,
)

sync_data_task = PythonOperator(
    task_id='sync_sample_data',
    python_callable=sync_sample_data,
    dag=dag,
)

generate_report_task = PythonOperator(
    task_id='generate_daily_report',
    python_callable=generate_daily_report,
    dag=dag,
)

quality_check_task = PythonOperator(
    task_id='data_quality_check',
    python_callable=data_quality_check,
    dag=dag,
)

optimize_task = PythonOperator(
    task_id='optimize_tables',
    python_callable=optimize_tables,
    dag=dag,
)

# 定义任务依赖关系
check_connection_task >> sync_data_task >> [generate_report_task, quality_check_task] >> optimize_task 